"""
词法分析器 V3.0
完全重新设计，支持完整的SQL语法
"""

import ply.lex as lex
import re
from typing import List, Dict, Any, Tuple

class LexerV3:
    """词法分析器 V3.0"""
    
    # 定义token类型（PLY要求作为类属性）
    tokens = [
        # 标识符和字面量
        'IDENTIFIER', 'STRING_LITERAL', 'NUMBER_LITERAL', 'BOOLEAN_LITERAL',
        
        # 操作符
        'PLUS', 'MINUS', 'MULTIPLY', 'DIVIDE', 'MODULO',
        'EQUALS', 'NOT_EQUALS', 'LESS_THAN', 'LESS_EQUAL', 'GREATER_THAN', 'GREATER_EQUAL',
        'ASSIGN',
        
        # 分隔符
        'COMMA', 'SEMICOLON', 'DOT', 'LPAREN', 'RPAREN',
        'LBRACKET', 'RBRACKET', 'LBRACE', 'RBRACE',
        
        # 其他
        'QUESTION_MARK',
        
        # SQL关键字
        'CREATE', 'TABLE', 'DROP', 'ALTER', 'INDEX',
        'SELECT', 'INSERT', 'INTO', 'VALUES', 'UPDATE', 'SET', 'DELETE', 'FROM', 'WHERE',
        'ORDER', 'BY', 'GROUP', 'HAVING', 'DISTINCT', 'AS',
        'INT', 'INTEGER', 'VARCHAR', 'CHAR', 'FLOAT', 'DOUBLE', 'BOOLEAN', 'BOOL',
        'DATE', 'TIME', 'DATETIME',
        'PRIMARY', 'KEY', 'FOREIGN', 'REFERENCES', 'UNIQUE', 'NOT', 'NULL', 'DEFAULT', 'AUTO_INCREMENT',
        'AND', 'OR', 'IN', 'EXISTS', 'BETWEEN', 'LIKE', 'IS',
        'COUNT', 'SUM', 'AVG', 'MIN', 'MAX',
        'ALL', 'ANY', 'SOME', 'CASE', 'WHEN', 'THEN', 'ELSE', 'END', 'LIMIT', 'OFFSET'
    ]
    
    def __init__(self):
        """初始化词法分析器"""
        self.lexer = None
        self.keywords = {
            # 数据定义语言 (DDL)
            'CREATE': 'CREATE',
            'TABLE': 'TABLE',
            'DROP': 'DROP',
            'ALTER': 'ALTER',
            'INDEX': 'INDEX',
            
            # 数据操作语言 (DML)
            'SELECT': 'SELECT',
            'INSERT': 'INSERT',
            'INTO': 'INTO',
            'VALUES': 'VALUES',
            'UPDATE': 'UPDATE',
            'SET': 'SET',
            'DELETE': 'DELETE',
            'FROM': 'FROM',
            'WHERE': 'WHERE',
            
            # 数据查询语言 (DQL)
            'ORDER': 'ORDER',
            'BY': 'BY',
            'GROUP': 'GROUP',
            'HAVING': 'HAVING',
            'DISTINCT': 'DISTINCT',
            'AS': 'AS',
            
            # 数据类型
            'INT': 'INT',
            'INTEGER': 'INTEGER',
            'VARCHAR': 'VARCHAR',
            'CHAR': 'CHAR',
            'FLOAT': 'FLOAT',
            'DOUBLE': 'DOUBLE',
            'BOOLEAN': 'BOOLEAN',
            'BOOL': 'BOOL',
            'DATE': 'DATE',
            'TIME': 'TIME',
            'DATETIME': 'DATETIME',
            
            # 约束
            'PRIMARY': 'PRIMARY',
            'KEY': 'KEY',
            'FOREIGN': 'FOREIGN',
            'REFERENCES': 'REFERENCES',
            'UNIQUE': 'UNIQUE',
            'NOT': 'NOT',
            'NULL': 'NULL',
            'DEFAULT': 'DEFAULT',
            'AUTO_INCREMENT': 'AUTO_INCREMENT',
            
            # 操作符和逻辑
            'AND': 'AND',
            'OR': 'OR',
            'IN': 'IN',
            'EXISTS': 'EXISTS',
            'BETWEEN': 'BETWEEN',
            'LIKE': 'LIKE',
            'IS': 'IS',
            
            # 聚合函数
            'COUNT': 'COUNT',
            'SUM': 'SUM',
            'AVG': 'AVG',
            'MIN': 'MIN',
            'MAX': 'MAX',
            
            # 其他
            'ALL': 'ALL',
            'ANY': 'ANY',
            'SOME': 'SOME',
            'CASE': 'CASE',
            'WHEN': 'WHEN',
            'THEN': 'THEN',
            'ELSE': 'ELSE',
            'END': 'END',
            'LIMIT': 'LIMIT',
            'OFFSET': 'OFFSET',

            # 复杂查询相关关键字
            'JOIN': 'JOIN',
            'INNER': 'INNER',
            'OUTER': 'OUTER',
            'LEFT': 'LEFT',
            'RIGHT': 'RIGHT',
            'FULL': 'FULL',
            'ON': 'ON',
            'UNION': 'UNION',
            'INTERSECT': 'INTERSECT',
            'EXCEPT': 'EXCEPT',
            'WITH': 'WITH',
            'RECURSIVE': 'RECURSIVE',
            'PARTITION': 'PARTITION',
            'OVER': 'OVER',
            'WINDOW': 'WINDOW',
            'ROWS': 'ROWS',
            'RANGE': 'RANGE',
            'PRECEDING': 'PRECEDING',
            'FOLLOWING': 'FOLLOWING',
            'CURRENT': 'CURRENT',
            'ROW': 'ROW',
            'UNBOUNDED': 'UNBOUNDED',
            'FILTER': 'FILTER',
            'WITHIN': 'WITHIN',
            'GROUP': 'GROUP',
            'ORDER': 'ORDER',
            'HAVING': 'HAVING',
        }
        
        # 构建词法分析器
        self.lexer = lex.lex(module=self)
    
    # 词法规则
    def t_IDENTIFIER(self, t):
        r'[a-zA-Z_][a-zA-Z0-9_]*'
        t.type = self.keywords.get(t.value.upper(), 'IDENTIFIER')
        return t
    
    def t_STRING_LITERAL(self, t):
        r"'([^'\\]|\\.)*'|\"([^\"\\]|\\.)*\""
        # 移除引号并处理转义字符
        value = t.value[1:-1]  # 移除首尾引号
        value = value.replace("\\'", "'").replace('\\"', '"')
        value = value.replace('\\n', '\n').replace('\\t', '\t').replace('\\r', '\r')
        t.value = value
        return t
    
    def t_NUMBER_LITERAL(self, t):
        r'\d+(\.\d+)?([eE][+-]?\d+)?'
        if '.' in t.value or 'e' in t.value.lower():
            t.value = float(t.value)
        else:
            t.value = int(t.value)
        return t
    
    def t_BOOLEAN_LITERAL(self, t):
        r'\b(TRUE|FALSE|true|false)\b'
        t.value = t.value.upper() == 'TRUE'
        t.type = 'BOOLEAN_LITERAL'
        return t
    
    # 操作符
    def t_PLUS(self, t):
        r'\+'
        return t
    
    def t_MINUS(self, t):
        r'-'
        return t
    
    def t_MULTIPLY(self, t):
        r'\*'
        # 检查是否在SELECT语句中，如果是则返回ASTERISK
        # 这里简化处理，直接返回MULTIPLY
        return t
    
    def t_DIVIDE(self, t):
        r'/'
        return t
    
    def t_MODULO(self, t):
        r'%'
        return t
    
    def t_EQUALS(self, t):
        r'='
        return t
    
    def t_NOT_EQUALS(self, t):
        r'!=|<>'
        return t
    
    def t_LESS_THAN(self, t):
        r'<'
        return t
    
    def t_LESS_EQUAL(self, t):
        r'<='
        return t
    
    def t_GREATER_THAN(self, t):
        r'>'
        return t
    
    def t_GREATER_EQUAL(self, t):
        r'>='
        return t
    
    def t_ASSIGN(self, t):
        r':='
        return t
    
    # 分隔符
    def t_COMMA(self, t):
        r','
        return t
    
    def t_SEMICOLON(self, t):
        r';'
        return t
    
    def t_DOT(self, t):
        r'\.'
        return t
    
    def t_LPAREN(self, t):
        r'\('
        return t
    
    def t_RPAREN(self, t):
        r'\)'
        return t
    
    def t_LBRACKET(self, t):
        r'\['
        return t
    
    def t_RBRACKET(self, t):
        r'\]'
        return t
    
    def t_LBRACE(self, t):
        r'\{'
        return t
    
    def t_RBRACE(self, t):
        r'\}'
        return t
    
    
    def t_QUESTION_MARK(self, t):
        r'\?'
        return t
    
    # 忽略空白字符
    def t_WHITESPACE(self, t):
        r'[ \t\n\r]+'
        pass
    
    # 注释处理
    def t_SINGLE_LINE_COMMENT(self, t):
        r'--.*'
        pass
    
    def t_MULTI_LINE_COMMENT(self, t):
        r'/\*.*?\*/'
        pass
    
    # 错误处理
    def t_error(self, t):
        """处理词法错误"""
        error_msg = f"词法错误: 非法字符 '{t.value[0]}' 在位置 {t.lexpos}"
        print(error_msg)
        t.lexer.skip(1)
        return t
    
    def tokenize(self, text: str) -> List[Dict[str, Any]]:
        """对输入文本进行词法分析"""
        self.lexer.input(text)
        tokens = []
        
        while True:
            tok = self.lexer.token()
            if not tok:
                break
            
            token_info = {
                'type': tok.type,
                'value': tok.value,
                'line': tok.lineno,
                'column': self._get_column(tok),
                'position': tok.lexpos
            }
            tokens.append(token_info)
        
        return tokens
    
    def _get_column(self, token) -> int:
        """获取token的列位置"""
        last_cr = self.lexer.lexdata.rfind('\n', 0, token.lexpos)
        if last_cr < 0:
            last_cr = -1
        return token.lexpos - last_cr
    
    def print_tokens(self, tokens: List[Dict[str, Any]]):
        """打印token流"""
        print("词法分析结果:")
        print("-" * 60)
        print(f"{'类型':<15} {'值':<20} {'行号':<6} {'列号':<6} {'位置':<6}")
        print("-" * 60)
        
        for token in tokens:
            value_str = str(token['value'])
            if len(value_str) > 18:
                value_str = value_str[:15] + "..."
            
            print(f"{token['type']:<15} {value_str:<20} {token['line']:<6} {token['column']:<6} {token['position']:<6}")
        
        print("-" * 60)
        print(f"总共识别到 {len(tokens)} 个token")

def main():
    """测试词法分析器"""
    lexer = LexerV3()
    
    # 测试SQL语句
    test_sql = """
    CREATE TABLE tst1 (id INT, name VARCHAR(50), age INT);
    INSERT INTO tst1 VALUES (1, 'Alice', 25);
    SELECT * FROM tst1 WHERE age > 20;
    DELETE FROM tst1 WHERE age > 20;
    """
    
    print("测试SQL语句:")
    print(test_sql)
    print("\n" + "="*80 + "\n")
    
    tokens = lexer.tokenize(test_sql)
    lexer.print_tokens(tokens)

if __name__ == "__main__":
    main()